package com.itcam.project.client;

import com.itcam.iahcommon.enums.ErrorCode;
import com.itcam.iahcommon.exception.BusinessException;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.PreDestroy;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

@Component
public class ArkClient {

    private final ArkService service;

    public ArkClient(@Value("${ark.api-key}") String apiKey,
                    @Value("${ark.base-url}") String baseUrl) {
        ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
        Dispatcher dispatcher = new Dispatcher();
        this.service = ArkService.builder()
                .dispatcher(dispatcher)
                .connectionPool(connectionPool)
                .baseUrl(baseUrl)
                .apiKey(apiKey)
                .build();
    }

    public String doStreamChat(String message) {
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content("你是人工智能助手.")
                .build());
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(message)
                .build());

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model("deepseek-v3-241226")
                .messages(messages)
                .build();

        StringBuilder contentBuilder = new StringBuilder();
        AtomicReference<Throwable> error = new AtomicReference<>();

        // 同步处理流式响应
        service.streamChatCompletion(request)
                .doOnError(error::set)
                .blockingForEach(choice -> {
                    if (!choice.getChoices().isEmpty()) {
                        contentBuilder.append(choice.getChoices().get(0).getMessage().getContent());
                    }
                });

        if (error.get() != null) {
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "AI流式响应错误");
        }

        return contentBuilder.toString();
    }

    @PreDestroy
    public void shutdown() {
        service.shutdownExecutor();
    }
}